import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset


# 定义全连接神经网络模型
class SheepModel(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(SheepModel, self).__init__()
        # 输入层到第一个隐藏层
        self.layers = nn.ModuleList()
        prev_size = input_size
        for hidden_size in hidden_sizes:
            self.layers.append(nn.Linear(prev_size, hidden_size))
            self.layers.append(nn.ReLU())
            prev_size = hidden_size
        # 最后一个隐藏层到输出层
        self.layers.append(nn.Linear(prev_size, output_size))

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


if __name__ == '__main__':
    # 生成一些随机数据用于示例
    # 输入特征数为10，样本数为100，标签数为1（回归问题）
    X = torch.randn(100, 10)
    y = torch.randn(100, 1)

    # 创建数据集和数据加载器
    dataset = TensorDataset(X, y)
    dataloader = DataLoader(dataset, batch_size=10, shuffle=True)



    # 设置超参数
    input_size = 10
    hidden_sizes = [32, 16]  # 你可以根据需要添加更多隐藏层及其大小
    output_size = 1
    learning_rate = 0.001
    num_epochs = 50

    # 实例化模型、损失函数和优化器
    model = FullyConnectedNN(input_size, hidden_sizes, output_size)
    criterion = nn.MSELoss()  # 均方误差损失，适用于回归问题
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 训练模型
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for inputs, targets in dataloader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(
            f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(dataloader):.4f}')

    # 评估模型（这里简单使用训练数据作为示例）
    model.eval()
    with torch.no_grad():
        total_loss = 0.0
        for inputs, targets in dataloader:
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            total_loss += loss.item()

        print(f'Average Loss on Training Data: {total_loss / len(dataloader):.4f}')
